#!/usr/bin/env python3
# encoding: utf-8
"""
@author: 孙伟伟
@contact: jakesun2020@163.com
@application:
@file: news.py
@time: 2021/11/19/019 17:15
"""
import torch
from tqdm import tqdm
# from torchtext.datasets import AG_NEWS
import os
# train_dataset, test_dataset = AG_NEWS(root="./data", split=("train","test"))
# loda_data_path = "./data"
# if not os.path.isdir(loda_data_path):
#     os.mkdir(loda_data_path)

# 选取torchtext中的文本分类数据集"AG_NEWS", 即新闻主题分类数据
# 将数值映射后的训练和验证集加载到内存中
# train_data, test_dataset = torchtext.datasets.AG_NEWS(root='./data/ag_news_csv/', split=('train', 'test'))
import pandas as pd

train_data1 = pd.read_csv("./data/ag_news_csv/AG_NEWS/train.csv")
mode_point = './model_point/'


label_map = {"1" : 0, "2": 1, "3": 2, "4" : 3}
'''
train.csv表示训练数据, 共12万条数据;   共由3列组成, 使用','进行分隔, 分别代表: 标签, 新闻标题, 新闻简述;    其中标签用"1", "2", "3", "4"表示
test.csv表示验证数据, 共7600条数据; 
里面有四个单词'World', 'Sports', 'Business', 'Sci/Tech'代表新闻的四个主题
'''

train_data1.columns = [0, 1, 2]
from torch.utils.data import Dataset,DataLoader
"""使用MyDataset类，构建自己的dataset
    Dataset:对于单个样本，从磁盘中读取训练数据集，特征，标签，稍微做一点数据的预处理，变成x,y的形成
    DataLoader：对于多个样本而言，把多个样本组合成一个mini_batch
"""

class MyDataset(Dataset):
    def __init__(self, data):
        '''
        处理数据的两种方法：
            结构化数据
            定义一个列表，把每个sample路径放到一个列表，标签放到另一个列表中
        :param filepath:
        '''
        xy = data
        self.len = len(xy)
        # self.text1 = xy.dataset.loc[:, 1]
        # self.text2 = xy.dataset.loc[:, 2]
        xy_ = data.dataset
        self.text = xy_[[1,2]].apply(lambda x: "".join(x), axis=1)
        self.lable = xy.dataset.loc[:, 0]

    def __getitem__(self, index):# 为了支持下标操作，即索引dataset[index]
        return self.text[index], self.lable[index]

    def __len__(self):
        return self.len

from torchtext.data.utils import get_tokenizer
from collections import Counter
from torchtext.vocab import Vocab

tokenizer = get_tokenizer('basic_english')      # 输入的字符串
counter = Counter()# 统计词频
for line in train_data1.loc[:, 2]:
    counter.update(tokenizer(line))# updata 如果要更新的关键字存在，对它的值进行求和操作，如果不存在，则添加，即对每个单词求词频
vocab = Vocab(counter)


# 准备数据处理管道
text_pipeline = lambda x: [vocab[token] for token in tokenizer(x)]      # token就是word，vocab[token]就是其对应的数字
'''lambda表达式的应用
    p = lambda x,y:x+y
    print(p(4,6))
    
'''


"""
    实现新闻主题分类任务，要完成五个步骤
        1.构建带有embedding层的文本分类模型
        2.对数据进行batch处理
        3.构建训练与验证函数
        4.进行模型训练和验证
        5.查看embedding层嵌入的词向量
"""
# 1.构建带有embedding层的文本分类模型
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence

BATCH_SIZE = 16

class TextSentiment(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_class):
        """
        文本分类模型
        :param vocab_size: 语料库中包含的不同词汇总数，即语料库的大小
        :param embed_dim:  词嵌入的大小,一个单词由多少个向量数值表示，即行向量的维数
        :param num_class:  分为的几类
        """
        super().__init__()
        # 实例化embedding层，传入3个参数，语料库大小，词嵌入维度，sparse=True表示每次对该层进行梯度求解时，只更新部分权重，
        self.embedding = nn.Embedding(vocab_size, embed_dim, sparse=True)# 输出为embed_dim列，几个单词就有几行，输出的维度与vocab_size无关
        """
            torch.nn.Embedding(num_embeddings, embedding_dim, padding_idx=None,
                                max_norm=None,  norm_type=2.0,   scale_grad_by_freq=False, 
                                sparse=False,  _weight=None)
                            num_embeddings: 词典大小，比如共有5000个词，那么输入就是5000
                            embedding_dim: 嵌入向量的维度，即用多少维来表示一个符号
                            padding_idx: 填充id, 比如输入长度为100，但是每个句子的长度不同，后面需要用统一的数字进行填充
                            max_norm: 最大范数，如果嵌入向量的范数超过了这个界限，就要进行归一化
        """
        self.fc = nn.Linear(embed_dim, num_class)
        self.init_weights()# 为各层初始化权重

    def init_weights(self):
        # 初始化权重的取值范围
        initrange = 0.5
        # 各层权重参数都是初始化为均匀分布
        self.embedding.weight.data.uniform_(-initrange, initrange)# 加上“_”表示就地操作，即把元数据直接改变了
        self.fc.weight.data.uniform_(-initrange, initrange)
        self.fc.bias.data.zero_()# 偏置初始化为0

    def forward(self, txt):
        """
        txt:表示文本进行数字化映射后的张量
        :param txt:文本数值映射为embedding的结果
        :return: 返回与类别尺寸相同的张量，用来判段分类类别
        """
        embedded = self.embedding(txt)# embeded.shape  (m, 32)  m是batchsize大小的数据中词汇总数 32表示词嵌入的长度
        # 将(m, 32)转化为(BATCH_SIZE, 32),便于通过fc层计算相应的损失

        return self.fc(embedded)# 减去新增的维度，然后转置回去馈送给fc层，预测出的结果为四个数字

# 实例化模型

VOCAB_SIZE = 958120# 语料库的大小，即包含多少不同的词汇数
# 指定词嵌入的维度
EMBED_DIM = 32
# 总的类别数
NUM_CLASS = 4
# 实例化模型
model = TextSentiment(VOCAB_SIZE, EMBED_DIM, NUM_CLASS)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 2.对数据进行batch处理，产生批次数据batch
def collate_fn1(batch):
    '''
    生成batch
    :param batch:由样本张量和对应标签的元组组成的batch_size大小的列表    [(label1, sample1), (lable2, sample2), ..., (labelN, sampleN)]
    :return:样本张量和标签各自的列表形式
    '''
    text_list = []
    text_max_len = 0
    tensor1 = torch.tensor([])
    for text in batch:
        text_list.append(text[0])

    for text in text_list:
        processed_text = torch.tensor(text_pipeline(text), dtype=torch.int64)
        if len(processed_text) > text_max_len:# 找到每个batch中的字符的最大长度
            text_max_len = len(processed_text)

    for text in text_list:# padding
        processed_text = torch.tensor(text_pipeline(text), dtype=torch.int64)
        if len(processed_text) <= text_max_len:
            pad = torch.zeros([text_max_len-len(processed_text)])
            text_tensor_ = torch.cat([processed_text, pad])
            tensor1 = torch.cat([tensor1, text_tensor_])
    # if len(tensor1) % 16 != 0:
    #     tensor1 = torch.tensor([])

    text_tensor2 = tensor1.reshape(16,-1).long()
    label_ = torch.tensor([label_list[1] for label_list in batch])
    label_ = label_ - torch.tensor([1])
    label = F.one_hot(label_, num_classes=NUM_CLASS)
    return text_tensor2, label


# 假设有一个输入
# batch = [(1, torch.tensor([3,23,2,8])),(0, torch.tensor([3,45,21,6])),(2, torch.tensor([30,145,1,66]))]
# res = collate_fn(batch)
# print(res)# (tensor([  3,  23,   2,   8,   3,  45,  21,   6,  30, 145,   1,  66]), tensor([1, 0, 2]))

# 3.构建训练和验证函数
from torch.utils.data import DataLoader
def train(train_data):
    # 训练损失和准确率都为0
    train_loss = 0
    tarin_acc = 0
    # 加载数据
    mydataset = MyDataset(train_data)
    data = DataLoader(dataset=mydataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0, collate_fn=collate_fn1)# collate_fn是对一个batch进行操作，例如padding

    train_dataloader = tqdm(data, desc="epoch:%d" % (epoch + 1), ncols=100)
    # data此时是一个一个batch的数据
    for i, (text, cls) in enumerate(train_dataloader):
        # 设置优化器初始梯度为0
        optimizer.zero_grad()
        text, cls = text.to(device), cls.to(device)
        output = model(text)# 一个batch的模型输出
        loss = criterion(output, cls)# 对真实标签和模型输出进行计算损失
        train_loss += loss.item()# 将该batch的损失加入到总损失中
        loss.backward()
        torch.nn.utils.clip_grad_norm(model.parameters(), 1)# 梯度剪枝，如果更新梯度小于设定阈值，则会将其限制在这个范围内，防止梯度爆炸
        #参数更新
        optimizer.step()
        # 准确率计算
        tarin_acc += (output.argmax(1) == cls).sum().item()# torch.argmax(input, dim=None)返回指定维度最大值的序号
        # if i == 1000:
        #     print("训练了1000个batch")
    # 调整优化器学习率
    scheduler.step()
    return train_loss / len(train_data), tarin_acc / len(train_data)

def valid(valid_data):
    '''验证函数'''
    loss = 0
    acc = 0
    mydataset = MyDataset(valid_data)
    data = DataLoader(dataset=mydataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0, collate_fn=collate_fn1)
    valid_dataloader = tqdm(data, desc="验证epoch:%d" % (epoch + 1), ncols=100)

    for text, cls in valid_dataloader:
        # 验证阶段，不需要求梯度
        with torch.no_grad():
            # 使用模型获得输出
            output = model(text)
            loss = criterion(output, cls)
            loss += loss.item()
            acc += (output.argmax(1) == cls).sum().item()
    return  loss / len(valid_data), acc / len(valid_data)

# 4.模型训练和验证
import time
# 随机划分数据集工具
from torch.utils.data.dataset import  random_split
# 指定训练轮次
N_EPOCHS = 10

# 定义初试的验证损失
min_valid_loss = float("inf")
# 损失函数选择
criterion = torch.nn.CrossEntropyLoss() # 0 1 2
criterion = criterion.to(device)
# 优化器选择
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
# 选择优化器步长调节方法StepLR，用来衰减学习率
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1, gamma=0.9)
'''
    torch.optim.lr_scheduler.StepLR(optimizer, step_size, gamma=0.1)
    等间隔调整学习率，调整倍数为 gamma 倍，调整间隔为 step_size。间隔单位是step。需要注意的是， step 通常是指 epoch
    step_size: 学习率下降间隔数，将学习率调整为 lr*gamma  
'''
# 从训练集中取出80%的数据作为训练集，另外20%作为验证集
train_len =  int(len(train_data1) * 0.8)

# 用random_split进行乱序划分，得到对应的训练集和验证集
sub_train_, sub_valid_ = random_split(train_data1, [train_len, len(train_data1) - train_len])

for epoch in range(N_EPOCHS):
    start_time = time.time()
    # 通过train和loss函数得到训练和验证的平均损失、平均准确率
    train_loss, train_acc = train(sub_train_)
    valid_loss, valid_acc = valid(sub_valid_)

    # 保存模型
    state = {"net":model.state_dict(), "optimizer":optimizer.state_dict(), 'epcoch':epoch+1}
    if not os.path.isdir(mode_point):
        os.mkdir(mode_point)
    save_path = mode_point + 'checkpoint_%d.cpt' % (epoch+1)
    torch.save(state, save_path)
    print("save_path:",save_path)

    # 计算训练和验证的总耗时
    secs = int(time.time() - start_time)
    mins = secs / 60# 分钟
    secs = secs %60# 秒

    # 打印训练和验证耗时，平均损失，平均准确率
    print('Epoch: %d' % (epoch + 1), " | time in %d minutes, %d seconds" % (mins, secs))
    print(f'\tLoss: {train_loss:.4f}(train)\t|\tAcc: {train_acc * 100:.1f}%(train)')
    print(f'\tLoss: {valid_loss:.4f}(valid)\t|\tAcc: {valid_acc * 100:.1f}%(valid)')
